;;;;;; attributes.scm - Class File Attributes module.    -*- Mode: Scheme -*-
;;;;;; Author: Erik Silkensen <silkense@colorado.edu>
;;;;;; Version: 19 Sep 2009

#lang scheme

(require "binary-io.scm")
(require "constant-pool.scm")

;;; This module provides constructors for Code, SourceFile, and
;;; LineNumberTable attributes. See Ch. 4.7 of the Java(TM) Virtual
;;; Machine Specification for documentation of these attributes:
;;; <http://java.sun.com/docs/books/jvms>

;;; Each attribtue obbject follows the pattern in Ch. 3 of Structure and 
;;; Interpretation of Computer Programs. The object is a procedure:
;;;     (let ((at (make-<attribute> <arg1>)))
;;;       ((at '<method>) <arg1> <arg2>))

;;; This module also provides procedures for several of the JVM
;;; instructions; each procedure generally writes the bytecode and then
;;; returns the integer offset the instruction had on the stack.

(provide
 aaload
 aastore
 aconst_null
 aload
 anewarray
 anewarray-class
 areturn
 arraylength
 astore
 checkcast
 dup
 getstatic
 goto
 iconst
 if_acmpeq
 if_acmpne
 iload
 invokespecial
 invokestatic
 invokevirtual
 ldc
 ldc-string
 make-code-attribute
 make-line-number-table-attribute
 make-source-file-attribute
 new
 new-class
 pc
 pop
 putstatic
 return
 tableswitch
 write-attributes)

(define (write-attributes attributes [out (current-output-port)])
  ;; Write a list of attributes to an output port. Attributes is a list
  ;; of objects; each object must have a write method (in the sense of
  ;; ((attribute 'write) out).
  (write-short (length attributes) out)
  (for-each (lambda (attribute)
              ((attribute 'write) out))
            attributes))

(define (make-code-attribute signature pool is-static)
  ;; Code attributes contain the JVM instructions and auxiliary
  ;; information for a method. Signature is the string descriptor of the
  ;; method that the code is for. Pool is the constant pool object for
  ;; the class that the code belongs to. Is-static is whether or not the
  ;; code is for a static method.
  (define (push value)
    ;; Adjust the depth of the operand stack, and update the max-stack
    ;; value. Return the updated depth of the operand stack.
    (set! stack (+ stack value))
    (when (> stack max-stack)
          (set! max-stack stack))
    stack)
  (define (reserve-locals locals)
    ;; Adjust the max-locals value such that at least a certain number
    ;; of local variables are allowed for in the code. Locals is the
    ;; number of local variable positions to reserve.
    (if (> locals max-locals)
        (set! max-locals locals)
        #t))
  (define (code-output)
    ;; Return the output port that bytecode is written to.
    code)
  (define (set-code-output! output)
    ;; Set the output port that bytecode is written to.
    (set! code output))
  (define (write [out (current-output-port)])
    ;; Write the Code attribute to an output port. See Ch. 4.7.3 of the
    ;; Java(TM) Virtual Machine Specification for more documentation.
    (let* ((code-bytes (get-output-bytes code))
           (code-length (bytes-length code-bytes))
           (att-length (+ 12 code-length)))
      (write-short name out)
      (write-int att-length out)
      (write-short max-stack out)
      (write-short max-locals out)
      (write-int code-length out)
      (write-bytes code-bytes out)
      (write-attributes '() out)    ;; exceptions
      (write-attributes '() out)))  ;; attributes
  (define code (open-output-bytes))
  (define stack 0)
  (define max-stack 0)
  (define max-locals (+ (- (sig-size signature #f)) (if is-static 0 1)))
  (define name ((pool 'Utf8) "Code"))
  (define this
    ;; The object is a procedure that returns a method from this table.
    (let ((method-table `((write . ,write)
                          (push . ,push)
                          (pool . ,pool)
                          (reserve-locals . ,reserve-locals)
                          (set-code-output! . ,set-code-output!)
                          (code-output . ,code-output))))
      (lambda (method)
        (let ((proc (assq method method-table)))
          (if proc
              (cdr proc)
              (error "undefined method -- CODE" method))))))
  this)

(define (make-line-number-table-attribute pool)
  ;; LineNumberTable attributes can optionally be used to map index
  ;; values in the code attribute to line numbers in the source file.
  (define (add-line! start-pc line-number)
    ;; Add a line into the Line Number Table. Start-pc is an index into
    ;; the code array at which the code for a new line begins. 
    ;; Line-number is the line in the source file corresponding to the
    ;; code for start-pc.
    (let ((line (cons start-pc line-number)))
      (set! lines (append lines (list line)))))
  (define (write [out (current-output-port)])
    ;; Write the Line Number Table to the output port. See Ch. 4.7.8 of
    ;; the Java(TM) Virtual Machine Specification.
    (let ((att-length (+ 2 (* 4 (length lines)))))
      (write-short name out)
      (write-int (+ 2 (* 4 (length lines))) out)
      (write-short (length lines) out)
      (for-each (lambda (line)
                  (write-short (car line) out)
                  (write-short (cdr line) out))
                lines)))
  (define lines '())
  (define name ((pool 'Utf8) "LineNumberTable"))
  (define this
    ;; The object is a procedure that returns a method from this table.
    (let ((method-table `((add-line! . ,add-line!)
                          (write . ,write))))
      (lambda (method)
        (let ((proc (assq method method-table)))
          (if proc
              (cdr proc)
              (error "undefined method -- LINE-NUMBER-TABLE" method))))))
  this)
(define (make-source-file-attribute source-file pool)
  ;; SourceFile attributes can optionally map generated Java class files
  ;; to their original source file.
  (define (write [out (current-output-port)])
    ;; Write the Source File attribute to the output port. See Ch. 4.7.7
    ;; of the Java(TM) Virtual Machine Specification.
    (write-short name out)
    (write-int 2)
    (write-short file-name))
  (define name ((pool 'Utf8) "SourceFile"))
  (define file-name ((pool 'Utf8) source-file))
  (define this
    ;; The object is a procedure that returns a method from this table.
    (let ((method-table `((write . ,write))))
      (lambda (method)
        (let ((proc (assq method method-table)))
          (if proc
              (cdr proc)
              (error "undefined method -- SOURCE-FILE" method))))))
  this)
(define (pc code)
  ;; Return the address of the next instruction to be written for a
  ;; code attribute. Code is a code attribute.
  (let ((code-output ((code 'code-output))))
    (bytes-length (get-output-bytes code-output))))

;;; The following procedures implement selected Java(TM) Virtual Machine
;;; instructions. In general, every instruction requires at least the
;;; code attribute object that the instruction is for. Each procedure
;;; will then adjust the necessary stack and local variable values in
;;; the code object.

(define (aaload code)
  ;; Load reference from array. 
  ;; Operand Stack: ...,arrayref,index => ...,value
  (let ((code-output ((code 'code-output))))
    (write-byte #x32 code-output)
    ((code 'push) 1)))

(define (aastore code)
  ;; Store reference into array.
  ;; Operand Stack: ...,arrayref,index,value => ...
  (let ((code-output ((code 'code-output))))
    (write-byte #x53 code-output)
    ((code 'push) -3)))

(define (aconst_null code)
  ;; Load null.
  ;; Operand Stack: ... => ...,null
  (let ((code-output ((code 'code-output))))
    (write-byte #x1 code-output)
    ((code 'push) 1)))

(define (aload index code)
  ;; Load reference from local variable. Index should be an index into
  ;; the local variable array of the current frame, where the value at
  ;; index is a reference.
  ;; Operand Stack: ... => ...,objectref
  (let ((code-output ((code 'code-output))))
    (if (and (>= index 0) (<= index 3))
        (write-byte (+ #x2a index) code-output)
        (begin (write-byte #x19 code-output)
               (write-byte index code-output)))
    ((code 'push) 1)))

(define (anewarray index code)
  ;; Create new array of reference. Index should be an index into the
  ;; constant pool where the value at index is a reference to a class,
  ;; array, or interface.
  ;; Operand Stack: ...,count => ...,arrayref
  (let ((code-output ((code 'code-output))))
    (write-byte #xbd code-output)
    (write-short index code-output)
    ((code 'push) 0)))

(define (anewarray-class class code)
  ;; Create a new array of a class. 
  ;; Operand Stack: ...,count => ...,arrayref
  (let ((pool (code 'pool)))
    (anewarray ((pool 'Class) class) code)))

(define (areturn code)
  ;; Return reference from method.
  ;; Operand Stack: ...,objectref => [empty]
  (let ((code-output ((code 'code-output))))
    (write-byte #xb0 code-output)
    ((code 'push) -1)))

(define (arraylength code)
  ;; Get length of array.
  ;; Operand Stack: ...,arrayref => ...,length
  (let ((code-output ((code 'code-output))))
    (write-byte #xbe code-output)
    ((code 'push) 0)))

(define (astore index code)
  ;; Store reference into local variable. Index should be an index into
  ;; the local variable array of the current frame.
  ;; Operand Stack: ...,objectref => ...
  (let ((code-output ((code 'code-output))))
    (if (and (>= index 0) (<= index 3))
        (write-byte (+ #x4b index) code-output)
        (begin (write-byte #x3a code-output)
               (write-byte index code-output)))
    ((code 'reserve-locals) (+ index 1))
    ((code 'push) -1)))

(define (checkcast class code)
  ;; Check whether object is of given type.
  ;; Operand Stack: ...,objectref => ...,objectref
  (let ((index (((code 'pool) 'Class) class))
        (code-output ((code 'code-output))))
    (write-byte #xc0 code-output)
    (write-short index code-output)
    ((code 'push) 0)))

(define (dup code)
  ;; Duplicate the top operand stack value. The value on top of the
  ;; operand stack must be a category 1 computational type (3.11.1).
  ;; Operand Stack: ...,value => ...,value,value
  (let ((code-output ((code 'code-output))))
    (write-byte #x59 code-output)
    ((code 'push) 1)))

(define (getstatic class name type code)
  ;; Get static field from class. Class should be the class containing
  ;; the static field. Name should be the name of the field, and type
  ;; should be the type descriptor of the field.
  ;; Operand Stack: ... => ...,value
  (let* ((pool (code 'pool))
         (index ((pool 'Fieldref) class name type))
         (size (hash-ref (pool 'sizes) index))
         (code-output ((code 'code-output))))
    (write-byte #xb2 code-output)
    (write-short index code-output)
    ((code 'push) size)))

(define (goto offset code)
  ;; Branch always. Jumps to the address of the goto opcode + offset.
  (let ((code-output ((code 'code-output))))
    (write-byte #xa7 code-output)
    (write-short offset code-output)
    ((code 'push) 0)))

(define (iconst value code)
  ;; Push integer constant. The actual opcode is selected based on the
  ;; value of the integer, and may not be iconst. The instructions that
  ;; are used are:
  ;;     -1 <= i <= 5       = iconst
  ;;   -126 <= i <= 127     = bipush
  ;; -32768 <= i <= 32767   = sipush
  ;; i < -32768 | i > 32767 = ldc
  ;; Operand Stack: ... => ...,value
  (let ((code-output ((code 'code-output))))
    (cond ((and (>= value -1) (<= value 5))
           (write-byte (+ #x3 value) code-output))
          ((and (>= value -126) (<= value 127))
           (write-byte #x10 code-output)
           (write-byte (if (< value 0) (+ value #x100) value) code-output))
          ((and (>= value -32768) (<= value 32767))
           (write-byte #x11 code-output)
           (write-short (if (< value 0) (+ value #x10000) value) code-output))
          (else
           (ldc (((code 'pool) 'Integer) value) code)))
    ((code 'push) 1)))

(define (if_acmpeq offset code)
  ;; Branch if reference comparison succeeds. If value1 = value2, then
  ;; execution continues at the address if_acmpeq opcode + offset. Else,
  ;; execution continues immediately after the if_acmpeq opcode.
  ;; Operand Stack: ...,value1,value2 => ...
  (if_acmp* #xa5 offset code))

(define (if_acmpne offset code)
  ;; Branch if reference comparison succeeds. If value1 = value2, then
  ;; execution continues immediately after the if_acmpne opcode. Else,
  ;; execution continues at the address if_acmpne opcode + offset.
  ;; Operand Stack: ...,value1,value2 => ...
  (if_acmp* #xa6 offset code))

(define (if_acmp* opcode offset code)
  (let ((code-output ((code 'code-output))))
    (write-byte opcode code-output)
    (write-short offset code-output)
    ((code 'push) -2)))

(define (iload index code)
  ;; Load int from local variable. Index should be an index into the
  ;; local variable array of the current frame, where the value at index
  ;; is an int.
  ;; Operand Stack: ... => ...,value
  (let ((code-output ((code 'code-output))))
    (if (and (>= index 0) (<= index 3))
        (write-byte (+ #x1a index) code-output)
        (begin (write-byte #x15 code-output)
               (write-byte index code-output)))
    ((code 'push) 1)))

(define (invokespecial class name type code)
  ;; Invoke instance method; special handling for superclass, private,
  ;; and instance initialization method invocations.
  ;; Operand Stack: ...,objectref,[arg1,[arg2,...]] => ...
  (invoke* #xb7 -1 class name type code))

(define (invokestatic class name type code)
  ;; Invoke a class (static) method.
  ;; Operand Stack: [arg1,[arg2,...]] => ...
  (invoke* #xb8 0 class name type code))

(define (invokevirtual class name type code)
  ;; Invoke instance method; dispatch based on class.
  ;; Operand Stack: ...,objectref,[arg1,[arg2,...]] => ...
  (invoke* #xb6 -1 class name type code))

(define (invoke* opcode push-offset class name type code)
  ;; Generic procedure that the different invoke procedures delegate to.
  ;; Opcode should be the opcode of the invoke instruction to write.
  ;; Push-offset is an integer i: size(methodref[class,name,type]) + i
  ;; is the effect of the invoke instruction on the operand stack.
  ;; Invokevirtual and invokespecial require a push-offset of -1 to
  ;; account for the extra reference used on the stack in addition to
  ;; the method size.
  (let* ((pool (code 'pool))
         (index ((pool 'Methodref) class name type))
         (size (hash-ref (pool 'sizes) index))
         (code-output ((code 'code-output))))
    (write-byte opcode code-output)
    (write-short index code-output)
    ((code 'push) (+ size push-offset))))

(define (ldc index code)
  ;; Push item from constant pool. Index is an index into the constant
  ;; pool. Uses ldc_w if index > 255.
  ;; Operand Stack: ... => ...,value
  (let ((code-output ((code 'code-output)))
        (opcode (if (<= index 255) #x12 #x13))
        (write-index (if (<= index 255) write-byte write-short)))
    (write-byte opcode code-output)
    (write-index index code-output)
    ((code 'push) 1)))

(define (ldc-string string code)
  ;; Push string from constant pool. String will be inserted into the
  ;; constant pool of the current class, if it is not there already.
  ;; Operand Stack: ... => ...,value
  (ldc (((code 'pool) 'String) string) code))

(define (new index code)
  ;; Create new object. Index should be an index into the constant pool
  ;; where the value at index is a reference to a class, array, or
  ;; interface type.
  ;; Operand Stack: ... => ...,objectref
  (let ((code-output ((code 'code-output))))
    (write-byte #xbb code-output)
    (write-short index code-output)
    ((code 'push) 1)))

(define (new-class class code)
  ;; Create new instance of a class. Class should be the name of the
  ;; class.
  ;; Operand Stack: ... => ...,objectref
  (let ((pool (code 'pool)))
    (new ((pool 'Class) class) code)))

(define (pop code)
  ;; Pop the top operand stack value.
  ;; Operand Stack: ...,value => ...
  (write-byte #x57 ((code 'code-output)))
  ((code 'push) -1))

(define (putstatic class name type code)
  ;; Set static field in class. Class should be the class containing the
  ;; static field. Name should be the name of the field, and type should
  ;; be the type descriptor of the field.
  ;; Operand Stack: ...,value => ...
  (let* ((pool (code 'pool))
         (index ((pool 'Fieldref) class name type))
         (size (hash-ref (pool 'sizes) index))
         (code-output ((code 'code-output))))
    (write-byte #xb3 code-output)
    (write-short index code-output)
    ((code 'push) (- size))))

(define (return code)
  ;; Return void from method.
  ;; ... => [empty]
  (write-byte #xb1 ((code 'code-output)))
  ((code 'push) 0))

(define (tableswitch padding default low high offsets code)
  ;; Access jump table by index and jump. Padding should be an integer 
  ;; 0 <= i <= 3 indicating the number of null bytes to use as padding.
  ;; [0]   <begin>
  ;; ...
  ;; [n]   tableswitch
  ;; [0-3] null bytes
  ;; [m]   ...
  ;; Padding should be chosen so that [m] is a multiple of 4 bytes from
  ;; the start of the current method. Default should be an integer i:
  ;; i + [n] is the address of the default code area of the switch. Low
  ;; and high should be the low and high index values for the switch,
  ;; inclusive. Offsets should be a list of target addresses for each of
  ;; the possible switch index values.
  ;; Operand Stack: ...,index => ...
  (let ((code-output ((code 'code-output))))
    (write-byte #xaa code-output)
    (cond ((= 3 padding)
           (write-short #x0 code-output)
           (write-byte #x0 code-output))
          ((= 2 padding)
           (write-short #x0 code-output))
          ((= 1 padding)
           (write-byte #x0 code-output)))
    (write-int default code-output)
    (write-int low code-output)
    (write-int high code-output)
    (for-each
     (lambda (offset)
       (write-int offset code-output))
     offsets))
  ((code 'push) -1))
